package com.ctvit.cch.core.service.impl;

import com.ctvit.cch.core.enums.MimeTypeEnum;
import com.ctvit.cch.core.service.IDocumentParseService;

import lombok.extern.slf4j.Slf4j;

import org.springframework.stereotype.Service;

import org.xml.sax.Attributes;
import org.xml.sax.ContentHandler;
import org.xml.sax.SAXException;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.io.FilenameUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.tika.exception.TikaException;
import org.apache.tika.extractor.EmbeddedDocumentExtractor;
import org.apache.tika.metadata.Metadata;
import org.apache.tika.parser.AutoDetectParser;
import org.apache.tika.parser.ParseContext;
import org.apache.tika.parser.Parser;
import org.apache.tika.parser.microsoft.OfficeParserConfig;
import org.apache.tika.parser.pdf.PDFParserConfig;
import org.apache.tika.sax.ContentHandlerDecorator;

import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * 文档解析服务实现类
 *
 * @author heyingcheng
 * @date 2025/11/13
 */
@Slf4j
@Service
public class DocumentParseServiceImpl implements IDocumentParseService {

    /**
     * 自定义ContentHandler用于插入Base64图片
     */
    private static class MarkdownImageHandler extends ContentHandlerDecorator {
        private final StringBuilder markdown = new StringBuilder();
        private final Map<String, String> imageMap;
        private String localName = null;

        public MarkdownImageHandler(Map<String, String> imageMap) {
            this.imageMap = imageMap;
        }

        @Override
        public void characters(char[] ch, int start, int length) {
            String text = new String(ch, start, length);
            if (this.localName != null) {
                if (this.localName.equals("h1")) {
                    markdown.append("# ").append(text);
                } else if (this.localName.equals("h2")) {
                    markdown.append("## ").append(text);
                } else if (this.localName.equals("h3")) {
                    markdown.append("### ").append(text);
                } else if (this.localName.equals("h4")) {
                    markdown.append("#### ").append(text);
                } else if (this.localName.equals("h5")) {
                    markdown.append("##### ").append(text);
                } else if (this.localName.equals("h6")) {
                    markdown.append("###### ").append(text);
                } else if (this.localName.equals("p")) {
                    markdown.append("\n").append(text);
                } else {
                    markdown.append(text);
                }
            } else {
                markdown.append(text);
            }
        }

        @Override
        public void startElement(String uri, String localName, String qName, Attributes attrs) {
            this.localName = localName;
            if ("img".equals(localName)) {
                String src = attrs.getValue("src");
                if (src != null && src.startsWith("embedded:")) {
                    String imageName = src.split(":")[1];
                    String extension = StringUtils.defaultIfBlank(FilenameUtils.getExtension(imageName), "png");
                    if (MimeTypeEnum.isSupportedImage(extension)) {
                        // 从Base64映射中获取图片数据
                        String base64Data = imageMap.get(imageName);
                        if (StringUtils.isNotBlank(base64Data)) {
                            // 构建Base64图片的Data URL
                            String mimeType = MimeTypeEnum.getMimeTypeByExtension(extension);
                            String dataUrl = "data:" + mimeType + ";base64," + base64Data;
                            markdown.append("![").append(imageName).append("](").append(dataUrl).append(")\n");
                        }
                    }
                }
            }
        }

        @Override
        public void endElement(String uri, String localName, String qName) {
            // 重置localName，避免影响后续文本处理
            if (this.localName != null && this.localName.equals(localName)) {
                this.localName = null;
            }
        }

        public String getMarkdown() {
            return markdown.toString();
        }
    }

    @Override
    public String extractText(InputStream inputStream) {
        // 初始化解析器、元数据和上下文
        Parser parser = new AutoDetectParser();
        Metadata metadata = new Metadata();
        ParseContext parseContext = new ParseContext();
        // 忽略 Office 文档中的页眉和页脚
        OfficeParserConfig officeParserConfig = new OfficeParserConfig();
        officeParserConfig.setIncludeHeadersAndFooters(false);
        parseContext.set(OfficeParserConfig.class, officeParserConfig);

        // PDF 配置，禁用 OCR(OCR 依赖 Tesseract 原生引擎, 需要额外依赖 tesseract 和 tess4j 性能开销大)
        PDFParserConfig pdfParserConfig = new PDFParserConfig();
        pdfParserConfig.setOcrStrategy(PDFParserConfig.OCR_STRATEGY.NO_OCR);
        parseContext.set(PDFParserConfig.class, pdfParserConfig);

        // 存储图片名称和Base64编码的映射
        Map<String, String> imageBase64Map = new LinkedHashMap<>();

        // 设置EmbeddedDocumentExtractor来处理嵌入的图片
        parseContext.set(EmbeddedDocumentExtractor.class, new EmbeddedDocumentExtractor() {
            @Override
            public boolean shouldParseEmbedded(Metadata metadata) {
                // 只处理图片类型
                return metadata.get(Metadata.CONTENT_TYPE) != null &&
                        metadata.get(Metadata.CONTENT_TYPE).startsWith("image/");
            }

            @Override
            public void parseEmbedded(InputStream inputStream, ContentHandler embeddedHandler, Metadata metadata,
                    boolean b) throws IOException, SAXException {
                String fileName = metadata.get("resourceName");
                if (StringUtils.isNotBlank(fileName)) {
                    // 将图片流转换为Base64字符串
                    byte[] imageBytes = inputStream.readAllBytes();
                    String base64String = Base64.encodeBase64String(imageBytes);
                    // 存储图片名称和Base64编码的映射
                    imageBase64Map.put(fileName, base64String);
                }
            }
        });

        // 创建MarkdownImageHandler实例并传入imageBase64Map
        MarkdownImageHandler contentHandler = new MarkdownImageHandler(imageBase64Map);

        // 开始解析文档
        try {
            parser.parse(inputStream, contentHandler, metadata, parseContext);
        } catch (IOException | SAXException | TikaException e) {
            throw new RuntimeException(e);
        }

        log.info("提取到的Markdown内容:[{}]", contentHandler.getMarkdown());

        return contentHandler.getMarkdown();
    }

}